import threading
import atexit  # 导入 atexit 模块

from sqlalchemy import create_engine, Column, Integer, String, DateTime, Index
from sqlalchemy.orm import sessionmaker, scoped_session
from sqlalchemy.ext.declarative import declarative_base
from logging_config import logger
from datetime import datetime

Base = declarative_base()
engine = create_engine('sqlite:///chat_history.db')
Session = scoped_session(sessionmaker(bind=engine))



class ChatMessage(Base):
    __tablename__ = 'chat_messages'
    id = Column(Integer, primary_key=True)
    sender_id = Column(String)
    sender_name = Column(String)
    message = Column(String)
    reply = Column(String)
    created_at = Column(DateTime, default=datetime.now)


Index('idx_sender_id', ChatMessage.sender_id)
Base.metadata.create_all(engine)  # 自动创建所有继承 Base 的模型对应的表


class DatabaseManager:
    # 批量保存消息的缓冲区
    message_buffer = []
    buffer_lock = threading.Lock()

    @staticmethod
    def save_message(sender_id, sender_name, message, reply):
        """保存消息到数据库，使用缓冲区批量保存"""
        with DatabaseManager.buffer_lock:
            DatabaseManager.message_buffer.append((sender_id, sender_name, message, reply))
            if len(DatabaseManager.message_buffer) >= 10:  # 达到 10 条消息时批量保存
                DatabaseManager._flush_buffer()

    @staticmethod
    def manual_flush():
        """手动刷新缓冲区"""
        DatabaseManager._flush_buffer()



    @staticmethod
    def _flush_buffer():
        """将缓冲区中的消息批量保存到数据库"""
        with DatabaseManager.buffer_lock:
            try:
                with Session() as session:
                    for sender_id, sender_name, message, reply in DatabaseManager.message_buffer:
                        chat_message = ChatMessage(
                            sender_id=sender_id,
                            sender_name=sender_name,
                            message=message,
                            reply=reply
                        )
                        session.add(chat_message)
                    session.commit()
                    DatabaseManager.message_buffer = []
            except Exception as e:
                logger.error(f"批量消息保存失败：{str(e)}")
                session.rollback()

    @staticmethod
    def get_recent_messages(sender_id, limit=10):
        """获取最近的消息记录"""
        try:
            with Session() as session:
                messages = session.query(ChatMessage).filter(
                    ChatMessage.sender_id == sender_id
                ).order_by(ChatMessage.id.desc()).limit(limit).all()
                return [{
                    "message": msg.message,
                    "reply": msg.reply,
                    "reply_time": msg.created_at.strftime("%Y-%m-%d %H:%M:%S")
                } for msg in reversed(messages)]  # 返回正序排列
        except Exception as e:
            logger.error(f"数据库查询失败：{str(e)}")
            return []


# 退出清理函数
def cleanup():
    DatabaseManager.manual_flush()
    Session.remove()    # 清理 scoped_session 的线程局部数据
    engine.dispose()

atexit.register(cleanup)